import os
from typing import Callable

import imageio
import jax
import jax.numpy as jnp
import numpy as np


def do_eval_single(
    env,
    task_id: int,
    agent,
    inference_horizon: int,
    rng,  # JAX random number generator
    is_render: bool = False,
    inference_fn: Callable = None,
):
    """Simulate a single evaluation run for given task_id"""
    list_trajs_act = []
    list_info = []
    list_frames = []
    # Reset environment
    obs, info = env.reset(options=dict(task_id=task_id, render_goal=is_render))
    # Reset agent's internal history for temporal ensembling
    if hasattr(agent, "reset_history"):
        agent.reset_history()
    goal = info["goal"]
    frame_goal = None
    if is_render:
        frame_goal = info["goal_rendered"]
    done = False
    # Simulation loop
    while not done:
        # Prepare observations for the agent
        obs_jnp = jnp.array(obs, dtype=jnp.float32)[jnp.newaxis, :]
        goal_jnp = jnp.array(goal, dtype=jnp.float32)[jnp.newaxis, :]

        # Evaluate the policy
        rng, infer_rng = jax.random.split(rng)

        # Use the provided inference function, with a fallback for compatibility
        call_inference_fn = (
            inference_fn if inference_fn is not None else agent.inference
        )
        tmp_trajs_act = call_inference_fn(
            rng=infer_rng,
            batch_obs={"proprio": obs_jnp},
            batch_goal={"proprio": goal_jnp},
        )
        list_trajs_act.append(np.array(tmp_trajs_act))
        # list_info.append(copy.deepcopy(agent._info)) # agent._info may not exist in jax agent

        for k in range(inference_horizon):
            action = np.array(tmp_trajs_act[0, k, :])
            obs, reward, terminated, truncated, info = env.step(action)
            if is_render:
                list_frames.append(env.render())
            done = terminated or truncated
            if done:
                break
    is_success = info["success"]
    return is_success, list_trajs_act, list_frames, frame_goal, list_info


def do_eval_stats(
    env,
    agent,
    num_eval_episode: int,
    inference_horizon: int,
    rng,
    inference_fn: Callable = None,
):
    list_is_success = []
    list_length = []
    for k in range(1, 6):
        for i in range(num_eval_episode):
            rng, eval_rng = jax.random.split(rng)
            is_success, list_trajs, list_frames, frame_goal, list_info = do_eval_single(
                env,
                k,
                agent,
                inference_horizon,
                eval_rng,
                is_render=False,
                inference_fn=inference_fn,
            )
            list_is_success.append(float(is_success))
            tmp_length = len(list_trajs) * inference_horizon
            if is_success:
                list_length.append(tmp_length)
            print("task_id", k, "run", i, "len", tmp_length, "is_success", is_success)
    overall_success = np.array(list_is_success).mean()
    success_length = np.array(list_length).mean() if len(list_length) > 0 else 0.0
    print("overall_success:", overall_success, "success_length:", success_length)
    return overall_success, success_length


def do_eval_and_render(
    env,
    agent,
    num_eval_episode: int,
    inference_horizon: int,
    rng,
    video_dir: str,
    inference_fn: Callable = None,
):
    """Run evaluation and render the episodes, saving them as videos."""
    for k in range(1, 6):  # Iterate through tasks
        for i in range(num_eval_episode):
            print(f"Rendering Task {k}, Episode {i + 1}/{num_eval_episode}...")
            rng, eval_rng = jax.random.split(rng)
            is_success, _, list_frames, _, _ = do_eval_single(
                env,
                k,
                agent,
                inference_horizon,
                eval_rng,
                is_render=True,
                inference_fn=inference_fn,
            )

            if not list_frames:
                print(
                    f"Warning: No frames rendered for Task {k}, Episode {i + 1}. Skipping video save."
                )
                continue

            # Save the video
            video_path = os.path.join(
                video_dir, f"task_{k}_episode_{i + 1}_success_{is_success}.mp4"
            )
            try:
                imageio.mimsave(video_path, list_frames, fps=30)
                print(f"Saved video to {video_path}")
            except Exception as e:
                print(f"Error saving video for Task {k}, Episode {i + 1}: {e}")
